/*
 * $Id$
 *
 *                This source code is part of
 *
 *                 G   R   O   M   A   C   S
 *
 *          GROningen MAchine for Chemical Simulations
 *
 *                        VERSION 3.0
 *
 * Copyright (c) 1991-2001
 * BIOSON Research Institute, Dept. of Biophysical Chemistry
 * University of Groningen, The Netherlands
 *
 * This program is free software; you can redistribute it and/or
 * modify it under the terms of the GNU General Public License
 * as published by the Free Software Foundation; either version 2
 * of the License, or (at your option) any later version.
 *
 * If you want to redistribute modifications, please consider that
 * scientific software is very special. Version control is crucial -
 * bugs must be traceable. We will be happy to consider code for
 * inclusion in the official distribution, but derived work must not
 * be called official GROMACS. Details are found in the README & COPYING
 * files - if they are missing, get the official version at www.gromacs.org.
 *
 * To help us fund GROMACS development, we humbly ask that you cite
 * the papers on the package - you can find them in the top README file.
 *
 * Do check out http://www.gromacs.org , or mail us at gromacs@gromacs.org .
 *
 * And Hey:
 * Gyas ROwers Mature At Cryogenic Speed
 */


/* This line is only for CVS version info */


#include "statutil.h"
#include "typedefs.h"
#include "smalloc.h"
#include "vec.h"
#include "copyrite.h"
#include "statutil.h"
#include "tpxio.h"
#include "fftgrid.h"
#include "index.h"
#include "pbc.h"
#include "get_charges.h"
#include "centerofmass.h"
#include "rdf_grid.h"
#include "get_potential.h"
#include "xvgr.h"
#include "get_xplor.h"
#include "rmpbc.h"

rvec ngrid={1,1,1};

real*** Rho_sol;
real*** Rho_slt;
real*** Rho_tot;
real*** Phi_sol;
real*** Phi_slt;
real*** Phi_tot;
real* Charges;
enum { euSel,euRect, euTric, euCompact, euNR};

static void put_molecule_com_in_box(int unitcell_enum,int ecenter,
                                    t_block *mols,
                                    int natoms,t_atom atom[],
                                    int ePBC,matrix box,rvec x[])
{
    atom_id i,j;
    int     d;
    rvec    com,new_com,shift,box_center;
    real    m;
    double  mtot;
    t_pbc   pbc;

    calc_box_center(ecenter,box,box_center);
    set_pbc(&pbc,ePBC,box);
    if (mols->nr <= 0)
        gmx_fatal(FARGS,"There are no molecule descriptions. I need a tpr file for this pbc option.");
    for(i=0; (i<mols->nr); i++) {
        /* calc COM */
        clear_rvec(com);
        mtot = 0;
        for(j=mols->index[i]; (j<mols->index[i+1] && j<natoms); j++) {
            m = atom[j].m;
            for(d=0; d<DIM; d++)
                com[d] += m*x[j][d];
            mtot += m;
        }
        /* calculate final COM */
        svmul(1.0/mtot, com, com);

        /* check if COM is outside box */
        copy_rvec(com,new_com);
        switch (unitcell_enum) {
        case euRect:
            put_atoms_in_box(box,1,&new_com);
            break;
        case euTric:
            put_atoms_in_triclinic_unitcell(ecenter,box,1,&new_com);
            break;
        case euCompact:
            put_atoms_in_compact_unitcell(ePBC,ecenter,box,1,&new_com);
            break;
        }
        rvec_sub(new_com,com,shift);
        if (norm2(shift) > 0) {
            if (debug)
                fprintf (debug,"\nShifting position of molecule %d "
                         "by %8.3f  %8.3f  %8.3f\n", i+1, PR_VEC(shift));
            for(j=mols->index[i]; (j<mols->index[i+1] && j<natoms); j++) {
                rvec_inc(x[j],shift);
            }
        }
    }
}

static void put_residue_com_in_box(int unitcell_enum,int ecenter,
                                   int natoms,t_atom atom[],
                                   int ePBC,matrix box,rvec x[])
{
    atom_id i, j, res_start, res_end, res_nat;
    int     d, presnr;
    real    m;
    double  mtot;
    rvec    box_center,com,new_com,shift;

    calc_box_center(ecenter,box,box_center);

    presnr = NOTSET;
    res_start = 0;
    clear_rvec(com);
    mtot = 0;
    for(i=0; i<natoms+1; i++) {
#ifdef OLDERVERSION
    	if (i == natoms || (presnr != atom[i].resnr && presnr != NOTSET)) {
#else
        if (i == natoms || (presnr != atom[i].resind && presnr != NOTSET)) {
#endif
        	/* calculate final COM */
            res_end = i;
            res_nat = res_end - res_start;
            svmul(1.0/mtot, com, com);

            /* check if COM is outside box */
            copy_rvec(com,new_com);
            switch (unitcell_enum) {
            case euRect:
                put_atoms_in_box(box,1,&new_com);
                break;
            case euTric:
                put_atoms_in_triclinic_unitcell(ecenter,box,1,&new_com);
                break;
            case euCompact:
                put_atoms_in_compact_unitcell(ePBC,ecenter,box,1,&new_com);
                break;
            }
            rvec_sub(new_com,com,shift);
            if (norm2(shift)) {
                if (debug)
#ifdef OLDERVERSION
                    fprintf (debug,"\nShifting position of residue %d (atoms %u-%u) "
                             "by %g,%g,%g\n", atom[res_start].resnr+1,
                             res_start+1, res_end+1, PR_VEC(shift));
#else
                    fprintf (debug,"\nShifting position of residue %d (atoms %u-%u) "
                             "by %g,%g,%g\n", atom[res_start].resind+1,
                             res_start+1, res_end+1, PR_VEC(shift));
#endif
                for(j=res_start; j<res_end; j++)
                    rvec_inc(x[j],shift);
            }
            clear_rvec(com);
            mtot = 0;

            /* remember start of new residue */
            res_start = i;
        }
        if (i < natoms) {
            /* calc COM */
            m = atom[i].m;
            for(d=0; d<DIM; d++)
                com[d] += m*x[i][d];
            mtot += m;
#ifdef OLDERVERSION
            presnr = atom[i].resnr;
#else
            presnr = atom[i].resind;
#endif
        }
    }
}

static void center_x(int ecenter,rvec x[],matrix box,int n,int nc,atom_id ci[]){
    int i,m,ai;
    rvec cmin,cmax,box_center,dx;

    if (nc > 0) {
        copy_rvec(x[ci[0]],cmin);
        copy_rvec(x[ci[0]],cmax);
        for(i=0; i<nc; i++) {
            ai=ci[i];
            for(m=0; m<DIM; m++) {
                if (x[ai][m] < cmin[m])
                    cmin[m]=x[ai][m];
                else if (x[ai][m] > cmax[m])
                    cmax[m]=x[ai][m];
            }
        }
        calc_box_center(ecenter,box,box_center);
        for(m=0; m<DIM; m++)
            dx[m] = box_center[m]-(cmin[m]+cmax[m])*0.5;

        for(i=0; i<n; i++)
            rvec_inc(x[i],dx);
    }
}

static void acc_x(int ncent, int *cindex, rvec *x_in, rvec *x_out){
	for(int ii=0;ii<ncent;ii++){
		int i=cindex[ii];
		x_out[i][XX]+=x_in[i][XX];
		x_out[i][YY]+=x_in[i][YY];
		x_out[i][ZZ]+=x_in[i][ZZ];
	}
}
static void avg_x(int ncent, int *cindex,int n, rvec *x){
	for(int ii=0;ii<ncent;ii++){
		int i=cindex[ii];
		x[i][XX]/=(real) n;
		x[i][YY]/=(real) n;
		x[i][ZZ]/=(real) n;
	}
}
static int calc_dipole(rvec com,rvec dip, int ncent, int * cindex, real *q, rvec *x){
	com[XX]=0.0;com[YY]=0.0;com[ZZ]=0.0;
	dip[XX]=0.0;dip[YY]=0.0;dip[ZZ]=0.0;
	for(int ii=0, i=cindex[ii];ii<ncent;ii++, i=cindex[ii]){
		com[XX]+=x[i][XX];
		com[YY]+=x[i][YY];
		com[ZZ]+=x[i][ZZ];
		dip[XX]+=q[i]*x[i][XX];
		dip[YY]+=q[i]*x[i][YY];
		dip[ZZ]+=q[i]*x[i][ZZ];
	}
	com[XX]/=(real) ncent;
	com[YY]/=(real) ncent;
	com[ZZ]/=(real) ncent;
	return 0;
}

void gmx_meancharges(int argc,char *argv[])
{
#ifdef OLDERVERSION
	static char *desc[] = {
#else
	const char *desc[] = {
#endif

	"this is a small test program meant to serve as a template ",
    "when writing your own analysis tools. The advantage of ",
    "using gromacs for this is that you have access to all ",
    "information in the topology, and your program will be ",
    "able to handle all types of coordinates and trajectory ",
    "files supported by gromacs. Go ahead and try it! ",
    "This test version just writes the coordinates of an ",
    "arbitrary atom to standard out for each frame. You can ",
    "select which atom you want to examine with the -n argument."
  };
  FILE *fp=NULL;
  FILE *fpxplor=NULL;
  t_topology top;
  int        ePBC;
  char       title[STRLEN], *grpnm;
  t_trxframe fr;
  rvec       *xtop;
  matrix     box;
  int        status,status_out=0;
  int        flags = TRX_READ_X;
  int        flagsb;
  int ncent;
  static rvec  trans = {0,0,0};
  static real filter=0.4;
#ifndef OLDVERSION
  output_env_t oenv;
#endif
  int nx,ny,nz;
  bool bIndex, bCenter=TRUE;
  atom_id      *cindex;
  bool         bRmPBC,bPBCWhole,bPBCcomRes,bPBCcomMol,bPBCcomAtom,bPBC,bNoJump,bCluster;

  bool 		bOXT=FALSE,bDipole=FALSE, bXplor=FALSE;
  bool 		bFilter=FALSE;
  bool         bTrans=FALSE;

  const char  *warn;
  const char     *distfile;
  t_atoms      *atoms=NULL;
  rvec *x;

  int pbc_enum;
  enum
  {
      epSel,
      epNone,
      epComMol,
      epComRes,
      epComAtom,
      epNojump,
      epCluster,
      epWhole,
      epNR
  };
#ifdef OLDERVERSION
  char *pbc_opt[epNR + 1] =
      { NULL, "none", "mol", "res", "atom", "nojump", "cluster", "whole",
          NULL };
  char *unitcell_opt[euNR+1] =
      { NULL, "rect", "tric", "compact", NULL };
  char *center_opt[ecNR+1] =
      { NULL, "tric", "rect", "zero", NULL };
#else
  const char *pbc_opt[epNR + 1] =
      { NULL, "none", "mol", "res", "atom", "nojump", "cluster", "whole",
          NULL };
  const char *unitcell_opt[euNR+1] =
      { NULL, "rect", "tric", "compact", NULL };
  const char *center_opt[ecNR+1] =
      { NULL, "tric", "rect", "zero", NULL };
#endif
  int unitcell_enum;
  enum
  { ecSel, ecTric, ecRect, ecZero, ecNR};
  int ecenter;

  t_pargs
		  pa[] ={
				  {"-grid", FALSE, etRVEC, {&ngrid}, "Grid size"},
                  { "-pbc", FALSE, etENUM,
                      { pbc_opt },
                      "PBC treatment (see help text for full description)" },
                  { "-ur", FALSE, etENUM,
                      { unitcell_opt }, "Unit-cell representation" },
                  { "-dipole", FALSE, etBOOL,
                      { &bDipole }, "The center molecule is a dipole" },
                  { "-filter", FALSE, etREAL,
                      {&filter}, "The center molecule is a dipole" },
                  { "-xplor", FALSE, etBOOL,
                      {&bXplor}, "Write a potential map in Xplor format" },
                    { "-trans", FALSE, etRVEC,
                        { trans },
                        "All coordinates will be translated by trans. This "
                        "can advantageously be combined with -pbc mol -ur "
                        "compact." },
                  { "-boxcenter", FALSE, etENUM,
                      { center_opt }, "Center for -pbc and -center" },
  };

  t_filenm fnm[] = {
    { efTPS,  NULL,  NULL, ffREAD },   /* this is for the topology */
    { efTRX, "-f", NULL, ffREAD },      /* and this for the trajectory */
    { efXVG, "-dist", "distr",    ffOPTWR  },
    { efTRX, "-oxt","coord.xtc", ffOPTWR },
    { efNDX, NULL,  NULL,         ffOPTRD },
  };
#define NFILE asize(fnm)
#define NPA asize(pa)


  CopyRight(stderr,argv[0]);


  /* This is the routine responsible for adding default options,
   * calling the X/motif interface, etc. */
#ifdef OLDERVERSION
  parse_common_args(&argc,argv,PCA_CAN_TIME | PCA_CAN_VIEW,
		    NFILE,fnm,asize(pa),pa,asize(desc),desc,0,NULL);
#else
  parse_common_args(&argc,argv,PCA_CAN_TIME | PCA_CAN_VIEW,
		    NFILE,fnm,asize(pa),pa,asize(desc),desc,0,NULL, &oenv);
#endif

  /* We don't need any topology information to write the coordinates,
   * but to show how it works we start by writing the name and
   * charge of the selected atom. It returns a boolean telling us
   * whether the topology was found and could be read
   */


    pbc_enum   = nenum(pbc_opt);
    bPBCWhole  = pbc_enum==epWhole;
    bPBCcomRes = pbc_enum==epComRes;
    bPBCcomMol = pbc_enum==epComMol;
    bPBCcomAtom= pbc_enum==epComAtom;
    bNoJump    = pbc_enum==epNojump;
    bCluster   = pbc_enum==epCluster;
    bPBC       = pbc_enum!=epNone;
    bTrans     = opt2parg_bSet("-trans",NPA,pa);
    bFilter    = opt2parg_bSet("-filter",NPA,pa);
    unitcell_enum = nenum(unitcell_opt);
    ecenter    = nenum(center_opt) - ecTric;

    /*
    int o;
    for(o=0;o<5;o++) printf("%s ",center_opt[o]);
    printf(" %d \n", nenum2(center_opt));
    printf(" ecenters %d %d \n",ecenter,ecTric);
    return;
     */

    bRmPBC = bPBCWhole || bPBCcomRes || bPBCcomMol;

    distfile = opt2fn_null("-dist",NFILE,fnm);
    bOXT = opt2bSet("-oxt",NFILE,fnm);

    if(distfile) fp = xvgropen(distfile,"","","",oenv);

	if (bOXT) {
		flagsb = flags | TRX_READ_X;
		status_out = open_trx(opt2fn("-oxt",NFILE,fnm),"w");
	}


    nx=(int) *ngrid;
    ny=(int) *(ngrid+1);
    nz=(int) *(ngrid+2);

    Rho_sol = mk_rgrid(nx, ny, nz);
    Rho_slt = mk_rgrid(nx, ny, nz);
    Rho_tot = mk_rgrid(nx, ny, nz);
	Phi_sol = mk_rgrid(nx, ny, nz);
	Phi_slt = mk_rgrid(nx, ny, nz);
	Phi_tot = mk_rgrid(nx, ny, nz);

	int i;

	real *ra=&Rho_sol[0][0][0];
	real *rb=&Rho_slt[0][0][0];
	real *rc=&Rho_tot[0][0][0];

	for (i = 0; i < nx*ny*nz; i++) {
		ra[i] = 0.0;
		rb[i] = 0.0;
		rc[i] = 0.0;
	}

	bIndex = ftp2bSet(efNDX, NFILE, fnm);
//	if (bIndex)
//		printf(" Success !!\n");

	read_tps_conf(ftp2fn(efTPS, NFILE, fnm), title, &top, &ePBC, &xtop, NULL,
			box, TRUE);
	sfree(xtop);
	atoms = &top.atoms;
	if(bOXT){
		int natom=atoms->nr;
		x=calloc(natom,sizeof(rvec));

	}
	if (bCenter) {
		if (bIndex) {
			printf("Select group for centering\n");
			rd_index(ftp2fn_null(efNDX, NFILE, fnm), 1, &ncent, &cindex, &grpnm);
		} else {
			gmx_fatal(FARGS, "Cannot do centering without Index file. Abort.");
		}
	}

//	printf(" %d\n",top.atoms.nr);
	Charges=malloc(sizeof(real)*top.atoms.nr);

	for(i=0;i<top.atoms.nr;i++) {
		Charges[i]=top.atoms.atom[i].q;
	}
  /* The first time we read data is a little special */
#ifdef OLDERVERSION
	read_first_frame(&status,ftp2fn(efTRX,NFILE,fnm),&fr,flags);
#else
	read_first_frame(oenv, &status,ftp2fn(efTRX,NFILE,fnm),&fr,flags);
#endif
	hist_t rsol, rslt, psol, pslt,ptot;

	if(distfile){
		if(setup_rdf_grid(&rsol)) return;
		if(setup_rdf_grid(&rslt)) return;
		if(setup_rdf_grid(&psol)) return;
		if(setup_rdf_grid(&pslt)) return;
		if(setup_rdf_grid(&ptot)) return;
	}
	/* This is the main loop over frames */
	int n=0;
	real Volume=0.0;

	rvec com, dip;
	rvec angle_loc;
	rvec vec_cell={0.0, 0.0, 0.0}, angle_cell={0.0, 0.0, 0.0};

	do {
		/* coordinates are available in the vector fr.x
		 * you can find this and all other structures in
		 * the types directory under the gromacs include dir.
		 * Note how flags determines whether to read x/v/f!
		 */

		printf(" %f %f %f \n ",fr.x[0][XX],fr.x[0][YY],fr.x[0][ZZ]);
		if (bCenter) center_x(ecenter,fr.x,fr.box,fr.natoms,ncent,cindex);
		t_pbc pbc;
		printf(" %f %f %f \n ",fr.x[0][XX],fr.x[0][YY],fr.x[0][ZZ]);

		set_pbc(&pbc,ePBC,fr.box);
		if (bRmPBC)
			rm_pbc(&(top.idef),ePBC,fr.natoms,fr.box,fr.x,fr.x);

		if (bPBCcomAtom) {
			switch (unitcell_enum) {
				case euRect:
				put_atoms_in_box(fr.box,fr.natoms,fr.x);
				break;
				case euTric:
				put_atoms_in_triclinic_unitcell(ecenter,fr.box,fr.natoms,fr.x);
				break;
				case euCompact:
				warn = put_atoms_in_compact_unitcell(ePBC,ecenter,fr.box,
						fr.natoms,fr.x);
				break;
			}
		}
		if (bPBCcomRes) {
			put_residue_com_in_box(unitcell_enum,ecenter,
					fr.natoms,atoms->atom,ePBC,fr.box,fr.x);
		}
		if (bPBCcomMol) {
			put_molecule_com_in_box(unitcell_enum,ecenter,
					&top.mols,
					fr.natoms,atoms->atom,ePBC,fr.box,fr.x);
		}

		if(bTrans) {
			for(int i=0; i<fr.natoms; i++)
				rvec_inc(fr.x[i],trans);
		}

		if (bOXT) {
			acc_x(ncent, cindex, fr.x, x);
		}
		if(get_charges(Rho_slt,Rho_sol,Charges,nx,ny,nz,ncent,cindex, &fr)) gmx_fatal(FARGS,"Cannot average charges. Abort!\n");
		for(int i=0;i<nx;i++)
			for(int j=0;j<ny;j++)
				for(int k=0;k<nz;k++)
					if(Rho_slt[i][j][k] != 0.0) {
						printf("ikka %d %d %d %f \n ",i,j,k,Rho_slt[i][j][k]);
					}


		exit(1);
		Volume+=det(fr.box);

		cell_angles(fr.box,&angle_loc[0], &angle_loc[1], &angle_loc[2]);
		for(int d=0;d<DIM; d++){
			vec_cell[d]+=norm(box[d]);
			angle_cell[d]+=angle_loc[d];
		}

		n++;

#ifdef OLDERVERSION
	} while(read_next_frame(status,&fr));
#else
	} while(read_next_frame(oenv,status,&fr));
#endif
	for(i=0;i<nx*ny*nz;i++){
		ra[i]=ra[i]/(real) n;
		rb[i]=rb[i]/(real) n;
	}

	matrix co,oc;
	for(int d=0;d<DIM; d++){
		vec_cell[d]  /=(real) n;
		angle_cell[d]/=(real) n;
	}
	angles_cell(co, vec_cell,angle_cell);
	if (bOXT) {
		avg_x(ncent, cindex, n, x);
#ifdef OLDERVERSION
		write_trx(status_out,ncent,cindex, &top.atoms,
              1,1.0,co,x,NULL);
#else
		write_trx(status_out,ncent,cindex, &top.atoms,
              1,1.0,co,x,NULL, NULL);
#endif
      }

	if(bDipole) calc_dipole(com, dip, ncent,cindex, Charges, fr.x);
	else gmx_calc_com(&top,x,ncent,cindex,com);

	transp_mat(co);
	calc_recipbox(co,oc);

	if(distfile) {
		do_rdf_grid(Rho_sol,nx,ny,nz,co,oc,com,&rsol);
		print_rdf_grid(fp,distfile,rsol,oenv);
		fprintf(fp,"&\n");

		do_rdf_grid(Rho_slt,nx,ny,nz,co,oc,com,&rslt);
		print_rdf_grid(fp,distfile,rslt,oenv);
		fprintf(fp,"&\n");
	}

	if(bFilter) macro_filter(nx,ny,nz, co, filter);
	Volume=Volume/(real) n;
	get_potential(Rho_sol, oc, nx,ny,nz, Volume, Phi_sol, bFilter);

	bFilter=FALSE;
	get_potential(Rho_slt, oc, nx,ny,nz, Volume, Phi_slt, bFilter);


	if(distfile) {
		if(bDipole) {
			do_rdf_dip_grid(Phi_sol,nx,ny,nz,co,oc,com,dip,&psol);
		} else do_rdf_grid(Phi_sol,nx,ny,nz,co,oc,com,&psol);


		print_rdf_grid(fp,distfile,psol,oenv);
		fprintf(fp,"&\n");

		if(bDipole) {
			do_rdf_dip_grid(Phi_slt,nx,ny,nz,co,oc,com,dip,&pslt);
		} else 	do_rdf_grid(Phi_slt,nx,ny,nz,co,oc,com,&pslt);
		print_rdf_grid(fp,distfile,pslt,oenv);
		fprintf(fp,"&\n");
	}

	ra=&Phi_sol[0][0][0];
	rb=&Phi_slt[0][0][0];
	rc=&Phi_tot[0][0][0];
	for(int i=0;i<nx*ny*nz;i++){
		rc[i]=ra[i]+rb[i];
	}

	if(distfile) {
		if(bDipole) {
			do_rdf_dip_grid(Phi_tot,nx,ny,nz,co,oc,com,dip,&ptot);
		} else do_rdf_grid(Phi_tot,nx,ny,nz,co,oc,com,&ptot);
		print_rdf_grid(fp,distfile,ptot,oenv);
	}


	if(bXplor) {
		fpxplor=fopen("Phi_slt.xplor","w+");
		get_xplor(Phi_slt,nx,ny,nz,co,fpxplor);
		ffclose(fpxplor);

		fpxplor=fopen("Phi_sol.xplor","w+");
		get_xplor(Phi_sol,nx,ny,nz,co,fpxplor);
		ffclose(fpxplor);


		fpxplor=fopen("Phi_tot.xplor","w+");
		get_xplor(Phi_tot,nx,ny,nz,co,fpxplor);
		ffclose(fpxplor);
	}


	if(distfile) ffclose(fp);

	thanx(stderr);

}

